(* $Id$ *)

(** A function [f] decides a predicate [p] if it returns [true] exactly
    when [p] holds and [false] exactly when [p] does not hold. *)
predicate decide (f, p) = 
    (forall x. (post (f) (x) (true)  <=> p (x)))
and (forall x. (post (f) (x) (false) <=> not (p (x))))
and (forall x. pre (f) (x)) 

(** The following will be imported from the standard library in the future. *)

(* [&] decides the boolean conjunction. *)
deferred val & : (bool * bool) -> bool
where (forall (x, y). pre (&) (x, y))
  and (forall (x, y). post (&) (x, y) (true) <=> (x = true and y = true))

(* [+'] computes the addition between two integers. *)
deferred val +? : (int * int) -> int
where (forall (x, y). pre (+?) (x, y))
  and (forall (x, y, z). post (+?) (x, y) (z) <=> z = x + y)

(* [-'] computes the substraction between two integers. *)
deferred val -? : (int * int) -> int
where (forall (x, y). pre (-?) (x, y))
  and (forall (x, y, z). post (-?) (x, y) (z) <=> z = x - y)

(* [max] is the logical function that returns the maximum of two integers. *)
deferred logic max  : (int * int) ~> int
where forall (x, y). (x < y => max (x, y) = y) and (y <= x => max (x, y) = x)

(* [max'] is the implementation of the [max] logical function. *)
deferred val max? : (int * int) -> int
where (forall (x, y). pre (max?) (x, y))
  and (forall (x, y, z). post (max?) (x, y) (z) <=> z = max (x, y))

(* The distance between two integers. *)
deferred logic dist : (int * int) ~> int
where forall (x, y). 
  ((x > y) => dist (x, y) = (x - y)) and 
  ((not (x > y)) => dist (x, y) = (y - x))

(* Comparison operators over integers. *)
deferred val <=?  : (int * int) -> bool where decide (<=?, <=)
deferred val >=?  : (int * int) -> bool where decide (>=?, >=)
deferred val >?   : (int * int) -> bool where decide (>?, >)
deferred val =?   : (int * int) -> bool
where (forall (x, y). pre ( =? ) (x, y))
  and (forall (x, y). post (=?) (x, y) (true) <=> x = y)
  and (forall (x, y). post (=?) (x, y) (false) <=> not (x = y))

(** Elements of the sets. *)

(* The type of the elements. *)
deferred type elt : *

(* Their equality is decidable. *)
deferred val =elt?  : (elt * elt) -> bool
where (forall x. pre (=elt?) (x))
  and (forall (x, y). post (=elt?) (x, y) (true) <=> x = y)
  and (forall (x, y). post (=elt?) (x, y) (false) <=> not (x = y))

(* There exists an order. *)
deferred predicate <elt : (elt * elt) ~> prop 
where (forall x. not (x <elt x))
  and (forall (x, y). (not (x <elt y) and not (x = y)) => y <elt x)
  and (forall (x, y, z). (x <elt y and y <elt z) => x <elt z)
  and (forall (x, y). (x <elt y) or (y <elt x) or (x = y))

(* This order is decidable. *)
deferred val <elt? : (elt * elt) -> bool where decide (<elt?, <elt)

(** The sets. *)

(* The type of an abstract set. *)
deferred type set : *

(* The membership predicate. *)
deferred predicate %in : (elt * set) ~> prop

(* The empty set. *)
deferred val empty : set
where forall x [x]. not (x %in empty)

deferred logic singleton : elt ~> set
where forall (x, y). x %in (singleton (y)) <=> (x = y)

deferred logic @ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 @ s2) <=> ((x %in s1) or (x %in s2))
deferred logic ^ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 ^ s2) <=> ((x %in s1) and (x %in s2))

deferred logic \\ : (set * set) ~> set
where forall (s1, s2, x). x %in (s1 \\ s2) <=> ((x %in s1) and not (x %in s2))

deferred predicate %C : (set * set) ~> prop
where forall (s1, s2). (s1 %C s2) <=> (forall x. (x %in s1) => (x %in s2))

lemma flatten_union: forall (x, s1, s2, s3). 
  x %in (s1 @ s2 @ s3) <=> (x %in s1 or x %in s2 or x %in s3)

lemma transitive_subset: forall (s1, s2, s3).
  (s1 %C s2 and s2 %C s3) => s1 %C s3

(* Inferior bound of a set. *)
predicate inf (x, s) = forall y. y %in s => x <elt y

lemma: forall (x, s1, s2). (inf (x, s1) and s2 %C s1) => inf (x, s2)
lemma: forall (x, s1, s2). (inf (x, s1 @ s2)) <=> (inf (x, s1) and inf (x, s2))
lemma: forall (s, x, y). (inf (x, s) and y <elt x) => inf (y, s)
lemma: forall (x, y). (inf (x, singleton (y)) <=> x <elt y)

(* Superior bound of a set. *)
predicate sup (x, s) = forall y. y %in s => y <elt x

lemma: forall (x, s1, s2). (sup (x, s1) and s2 %C s1) => sup (x, s2)
lemma: forall (x, s1, s2). (sup (x, s1 @ s2)) <=> (sup (x, s1) and sup (x, s2))
lemma: forall (s, x, y). (sup (x, s) and x <elt y) => sup (y, s)
lemma: forall (x, y). (sup (x, singleton (y)) <=> y <elt x)

(* Extensional equality. *)
predicate === (s1, s2) = forall x. (x %in s1) <=> (x %in s2)

(* Extensionality axiom. *)
axiom: forall (s1, s2). (s1 === s2) <=> (s1 = s2) 

(** Trees. *)

(* Inductive type for tree. *)
type tree =
  | Empty : tree
  | Node  : (int * tree * elt * tree) -> tree

(* The abstraction of trees to sets. *)
logic elements (t) = 
  | Empty -> empty
  | Node (h, l, x, r) -> elements (l) @ singleton (x) @ elements (r)

lemma root_in_elements: forall (h, l, x, r). x %in elements (Node (h, l, x, r))
lemma subl_in_elements: forall (h, l, x, r). 
  elements (l) %C elements (Node (h, l, x, r))
lemma subr_in_elements: forall (h, l, x, r). 
  elements (r) %C elements (Node (h, l, x, r))
lemma subl2_in_elements: forall (h, hl, ll, lx, lr, x, r). 
    elements (ll) %C elements (Node (h, Node (hl, ll, lx, lr), x, r))
and elements (lr) %C elements (Node (h, Node (hl, ll, lx, lr), x, r))
lemma subr2_in_elements: forall (h, l, x, hr, rl, rx, rr). 
  elements (rl) %C elements (Node (h, l, x, Node (hr, rl, rx, rr)))
and elements (rr) %C elements (Node (h, l, x, Node (hr, rl, rx, rr)))

(* A simple function: the emptyness test. *)
let is_empty t returns b where ((b = true) <=> (elements (t) === empty)) =
  match t with
    | Empty             -> true
    | Node (h, l, x, r) -> false
 end

(* Binary search trees. *)

inductive bst : tree ~> prop =
  | bst (Empty)
  | forall (h, l, x, r). 
      bst (l) and bst (r) and sup (x, elements (l)) and inf (x, elements (r))
      => bst (Node (h, l, x, r))

predicate bst_node (l, x, r) = 
   bst (l) and bst (r) and sup (x, elements (l)) and inf (x, elements (r))

lemma: forall (l, x, h, rl, rx, rr).
   bst_node (l, x, Node (h, rl, rx, rr)) =>
   sup (rx, elements (l) @ singleton (x) @ elements (rl))

lemma: forall (h, ll, lx, lr, x, r).
   bst_node (Node (h, ll, lx, lr), x, r) =>
   inf (lx, elements (lr) @ singleton (x) @ elements (r))

lemma: forall (lh, ll, lx, lrh, lrl, lrx, lrr, x, r).
   bst_node (Node (lh, ll, lx, Node (lrh, lrl, lrx, lrr)), x, r) =>
   inf (lrx, elements (lrr) @ singleton (x) @ elements (r))

lemma: forall (h, l, x, rh, rlh, rll, rlx, rlr, rx, rr).
   bst_node (l, x, Node (rh, Node (rlh, rll, rlx, rlr), rx, rr)) =>
   sup (rlx, elements (l) @ singleton (x) @ elements (rll))

(* Membership decision procedure over binary search trees. *)

let rec mem_bst (t, x) where bst (t) 
returns b where ((b = true) <=> (x %in elements (t))) =

    match t with

      | Empty ->
          false 
	    
      | Node (h, l, y, r) ->
          if (x =elt? y) then
            true
          else if (x <elt? y) then
            mem_bst (l, x)
          else 
            mem_bst (r, x)
  
 end

(* Height. *)

logic height t =
  | Empty -> 0
  | Node (h, l, x, r) -> 1 + (max (height (l), height (r))) 

lemma height_positive: forall t. height (t) >= 0

(* A tree is well-formed if the integer stored in its root node 
   is its height. *)

inductive wf : tree ~> prop =
  | wf (Empty)
  | forall (h, l, x, r). h = height (Node (h, l, x, r)) and wf (l) and wf (r)
    => wf (Node (h, l, x, r))

(* We can extract the height of a well-formed tree in constant time. *)

let height? (t) where wf (t) returns h where (h = height (t)) =
  match t with
    | Empty -> 0
    | Node (h, l, x, r) -> h
  end

(* The definition of AVL trees. *)

predicate sim_depth (l, r) = 2 >= dist (height (l), height (r)) 

predicate almost_sim_depth (l, r) = 3 >= dist (height (l), height (r)) 

inductive avl : tree ~> prop = 
  | avl (Empty)
  | forall (h, x, l, r).
      avl (l) and avl (r) and sim_depth (l, r)
    => avl (Node (h, l, x, r))

predicate avl_node (l, r) = 
  avl (l) and avl (r) and sim_depth (l, r) 

predicate almost_avl_node (l, r) = 
  avl (l) and avl (r) and almost_sim_depth (l, r) 

predicate wf_bst_avl (t) = 
  wf (t) and bst (t) and avl (t) 

(* Smart constructor of BST and AVL well-formed trees. *)

let mk_node (l, x, r) 
where bst_node (l, x, r) and avl_node (l, r) and wf (l) and wf (r)
returns t 
where (wf_bst_avl (t) and bst_node (l, x, r)
  and height (t) = max (height (l), height (r)) + 1 
  and elements (t) === (elements (l) @ singleton (x) @ elements (r)) 
  and t = Node (max (height (l), height (r)) + 1, l, x, r)) =
  Node (max? (height? (l), height? (r)) +? 1, l, x, r)

(* Balancing. *)
lemma union_assoc : forall (s1, s2, s3). (s1 @ (s2 @ s3)) === ((s1 @ s2) @ s3)

let bal (l, x, r) 
where wf (l) and wf (r) and almost_avl_node (l, r) and bst_node (l, x, r)
returns t 
where (wf (t) and avl (t) 
  and elements (t) === (elements (l) @ singleton (x) @ elements (r))
  and (avl_node (l, r) => height (t) = max (height (l), height (r)) + 1)
  and dist (1 + max (height (l), height (r)), height (t)) <= 1) =
  
  let hl = height? (l) in
  let hr = height? (r) in 

    if (hl -? hr) =? 3 then let Node (h, ll, lx, lr) = l in
	
        if height? (lr) <=? height? (ll) then 
	  mk_node (ll, lx, mk_node (lr, x, r))
        
        else let Node (h, lrl, lrx, lrr) = lr in 
	  assert bst_node (ll, lx, lrl) in
          mk_node (mk_node (ll, lx, lrl), lrx, mk_node (lrr, x, r))

    else if (hr -? hl) =? 3 then let Node (h, rl, rx, rr) = r in
	
        if height? (rl) <=? height? (rr) then 
	  mk_node (mk_node (l, x, rl), rx, rr)
        
        else let Node (h, rll, rlx, rlr) = rl in 
          mk_node (mk_node (l, x, rll), rlx, mk_node (rlr, rx, rr))

    else 
      mk_node (l, x, r)
